/**
 * @file    px4_estimator
 * @brief   subscribe odometry from sensors & publish pose to mavros for odometry - imu fusion
 * @author  FLAG Lab, BIT
 * @version 1.1 remove some shit
 * @date    2024-05-26
 */

#include <fsm_ctrl/px4_estimator.hpp>

using namespace std;



/*--------- variable ---------*/
/* mocap */
static bool is_mocap_init = false;
static Eigen::Vector3d pos_mocap;                   
static Eigen::Quaterniond quat_mocap;
static Eigen::Vector3d euler_mocap;
static Eigen::Vector3d pos_mocap_init;                   
static Eigen::Quaterniond quat_mocap_init;

/* lidar */
static Eigen::Vector3d pos_lidar;
static Eigen::Vector3d vel_lidar;                     
static Eigen::Quaterniond quat_lidar;
static Eigen::Vector3d euler_lidar; 

/* camera */
static Eigen::Vector3d pos_camera;
static Eigen::Vector3d vel_camera;                     
static Eigen::Quaterniond quat_camera;
static Eigen::Vector3d euler_camera;

/* FCU */
static Eigen::Vector3d pos_fcu;
static Eigen::Vector3d vel_fcu;                     
static Eigen::Quaterniond quat_fcu;
static Eigen::Vector3d euler_fcu;

/* IMU */
static Eigen::Vector3d acc_imu;                          
static Eigen::Quaterniond quat_imu;
static Eigen::Vector3d euler_imu;
static Eigen::Vector3d rate_imu;

/* TFmini */
static double tfmini_raw;
static double tfmini_map;

/* EKF */
static bool is_source_new = false;
static int flag_vision_source;
std_msgs::Bool ready;
geometry_msgs::PoseStamped vision;



/*--------- callback function ---------*/
/**
 * @brief  motion capture pose subscriber callback
 * @param  msg: geometry_msgs::PoseStamped from vrpn_client_ros
 * @return NULL
 */
void Mocap_Callback(const geometry_msgs::PoseStamped::ConstPtr &msg)
{
    int flag_mocap_frame = 1;         // optitrack frame in 104, 0: Z-up; 1: Y-up
    Eigen::Quaterniond pos_mocap_raw, quat_mocap_raw;

    if(flag_mocap_frame == 0)
    {
        if(!is_mocap_init)
        {
            pos_mocap_init = Eigen::Vector3d(msg->pose.position.x, msg->pose.position.y, msg->pose.position.z);
            quat_mocap_init = Eigen::Quaterniond(msg->pose.orientation.w, msg->pose.orientation.x, msg->pose.orientation.y, msg->pose.orientation.z);
            is_mocap_init = true;
        }
        pos_mocap_raw = Eigen::Quaterniond(0.0, msg->pose.position.x - pos_mocap_init[0], msg->pose.position.y - pos_mocap_init[1], msg->pose.position.z - pos_mocap_init[2]);
        quat_mocap_raw = Eigen::Quaterniond(msg->pose.orientation.w, msg->pose.orientation.x, msg->pose.orientation.y, msg->pose.orientation.z);
    }

    else if(flag_mocap_frame == 1)
    {
        if(!is_mocap_init)
        {
            pos_mocap_init = Eigen::Vector3d(msg->pose.position.x, -msg->pose.position.z, msg->pose.position.y);
            quat_mocap_init = Eigen::Quaterniond(msg->pose.orientation.w, msg->pose.orientation.x, msg->pose.orientation.z, msg->pose.orientation.y);
            is_mocap_init = true;
        }
        // pos_mocap_raw = Eigen::Quaterniond(0.0, msg->pose.position.x - pos_mocap_init[0], -msg->pose.position.z - pos_mocap_init[1], msg->pose.position.y - pos_mocap_init[2]);
        // quat_mocap_raw = Eigen::Quaterniond(msg->pose.orientation.w, msg->pose.orientation.x, msg->pose.orientation.z, msg->pose.orientation.y);
    }

    else {ROS_ERROR("Invalid Mocap Frame!!!");}

    // pos_mocap_raw = quat_mocap_raw * pos_mocap_raw * quat_mocap_raw.inverse();
    pos_mocap = Eigen::Vector3d(msg->pose.position.x, -msg->pose.position.z, msg->pose.position.y) - pos_mocap_init; 
    quat_mocap = Eigen::Quaterniond(msg->pose.orientation.w, msg->pose.orientation.x, msg->pose.orientation.z, msg->pose.orientation.y);
    euler_mocap = QuatToEuler(quat_mocap);

    if(flag_vision_source == 0)
    {
        vision.header.stamp = msg->header.stamp;
        vision.pose.position.x = pos_mocap[0] ;
        vision.pose.position.y = pos_mocap[1] ;
        vision.pose.position.z = pos_mocap[2] ;
        vision.pose.orientation.w = quat_mocap.w();
        vision.pose.orientation.x = quat_mocap.x();
        vision.pose.orientation.y = quat_mocap.y();
        vision.pose.orientation.z = quat_mocap.z();
        is_source_new = true;
    }
}


/**
 * @brief  lidar odometry subscriber callback
 * @param  msg: nav_msgs::Odometry from mid360
 * @return NULL
 */
void Lidar_Callback(const nav_msgs::Odometry::ConstPtr &msg)
{
    pos_lidar = Eigen::Vector3d(msg->pose.pose.position.x, msg->pose.pose.position.y, msg->pose.pose.position.z);
    vel_lidar = Eigen::Vector3d(msg->twist.twist.linear.x, msg->twist.twist.linear.y, msg->twist.twist.linear.z);
    quat_lidar = Eigen::Quaterniond(msg->pose.pose.orientation.w, msg->pose.pose.orientation.x, msg->pose.pose.orientation.y, msg->pose.pose.orientation.z);
    euler_lidar = QuatToEuler(quat_lidar);

    if(flag_vision_source == 1)
    {
        vision.header.stamp = msg->header.stamp;
        vision.pose.position.x = pos_lidar[0] ;
        vision.pose.position.y = pos_lidar[1] ;
        vision.pose.position.z = pos_lidar[2] ;
        vision.pose.orientation.w = quat_lidar.w();
        vision.pose.orientation.x = quat_lidar.x();
        vision.pose.orientation.y = quat_lidar.y();
        vision.pose.orientation.z = quat_lidar.z();
        is_source_new = true;
    }
}


/**
 * @brief  camera odometry subscriber callback
 * @param  msg: nav_msgs::Odometry from D435i
 * @return NULL
 */
void Camera_Callback(const nav_msgs::Odometry::ConstPtr &msg)
{
    pos_camera = Eigen::Vector3d(msg->pose.pose.position.x, msg->pose.pose.position.y, msg->pose.pose.position.z);
    vel_camera = Eigen::Vector3d(msg->twist.twist.linear.x, msg->twist.twist.linear.y, msg->twist.twist.linear.z);
    quat_camera = Eigen::Quaterniond(msg->pose.pose.orientation.w, msg->pose.pose.orientation.x, msg->pose.pose.orientation.y, msg->pose.pose.orientation.z);
    euler_camera = QuatToEuler(quat_camera);

    if(flag_vision_source == 2)
    {
        vision.header.stamp = msg->header.stamp;
        vision.pose.position.x = pos_camera[0] ;
        vision.pose.position.y = pos_camera[1] ;
        vision.pose.position.z = pos_camera[2] ;
        vision.pose.orientation.w = quat_lidar.w();
        vision.pose.orientation.x = quat_lidar.x();
        vision.pose.orientation.y = quat_lidar.y();
        vision.pose.orientation.z = quat_lidar.z();
        is_source_new = true;
    }
}


/**
 * @brief  fusion pose subscriber callback
 * @param  msg: geometry_msgs::PoseStamped from FCU
 * @return NULL
 */
void Pose_Callback(const geometry_msgs::PoseStamped::ConstPtr &msg)
{
    pos_fcu = Eigen::Vector3d(msg->pose.position.x, msg->pose.position.y, msg->pose.position.z);
    quat_fcu = Eigen::Quaterniond(msg->pose.orientation.w, msg->pose.orientation.x, msg->pose.orientation.y, msg->pose.orientation.z);
    euler_fcu = QuatToEuler(quat_fcu);
}


/**
 * @brief  fusion velocity subscriber callback
 * @param  msg: geometry_msgs::TwistStamped from FCU
 * @return NULL
 */
void Vel_Callback(const geometry_msgs::TwistStamped::ConstPtr &msg)
{
    vel_fcu = Eigen::Vector3d(msg->twist.linear.x, msg->twist.linear.y, msg->twist.linear.z);
}


/**
 * @brief  imu data subscriber callback
 * @param  msg: sensor_msgs::Imu from IMU of FCU
 * @return NULL
 */
void IMU_Callback(const sensor_msgs::Imu::ConstPtr &msg)
{
    acc_imu = Eigen::Vector3d(msg->linear_acceleration.x, msg->linear_acceleration.y, msg->linear_acceleration.z);
    quat_imu = Eigen::Quaterniond(msg->orientation.w, msg->orientation.x, msg->orientation.y, msg->orientation.z);
    euler_imu = QuatToEuler(quat_imu);
    rate_imu = Eigen::Vector3d(msg->angular_velocity.x, msg->angular_velocity.y, msg->angular_velocity.z);   
}


/**
 * @brief  map ranger data to height 
 * @param  range
 * @return NULL
 */
double MapHeight(double _range, double _roll, double _pitch)
{
    double height = _range*cos(atan(sqrt(tan(_roll)*tan(_roll) + tan(_pitch)*tan(_pitch) )));
    return height;
}

/**
 * @brief  ranger data subscriber callback
 * @param  msg: sensor_msgs::Range from TFmini
 * @return NULL
 */
void Ranger_Callback(const sensor_msgs::Range::ConstPtr &msg)
{
    tfmini_raw = msg->range;
    tfmini_map = MapHeight(tfmini_raw, euler_fcu(0), euler_fcu(1));
}



int main(int argc, char **argv)
{
    ros::init(argc, argv, "px4_estimator");
    ros::NodeHandle nh("~");

    /* parameter */
    nh.param("vision_source", flag_vision_source, 0);

    /* publisher */
    ros::Publisher fusion_ready_pub = nh.advertise<std_msgs::Bool>
        ("/fsm_ctrl/ekf_fusion_ready", 10);
    ros::Publisher vision_pub = nh.advertise<geometry_msgs::PoseStamped>
        ("/mavros/vision_pose/pose", 10);
    ros::Publisher debug_pub = nh.advertise<geometry_msgs::PoseStamped>
        ("/fsm_ctrl/estimator_debug", 10);
    
    /* subscriber */
    ros::Subscriber mocap_sub = nh.subscribe<geometry_msgs::PoseStamped>
        ("pose", 10, Mocap_Callback);
    ros::Subscriber lidar_sub = nh.subscribe<nav_msgs::Odometry>
        ("odom", 10, Lidar_Callback);    
    ros::Subscriber camera_sub = nh.subscribe<nav_msgs::Odometry>
        ("odom", 10, Camera_Callback);
    ros::Subscriber ranger_sub = nh.subscribe<sensor_msgs::Range>
        ("/tfmini", 10, Ranger_Callback);
    ros::Subscriber pose_sub = nh.subscribe<geometry_msgs::PoseStamped>
        ("/mavros/local_position/pose", 10, Pose_Callback);
    ros::Subscriber vel_sub = nh.subscribe<geometry_msgs::TwistStamped>
        ("/mavros/local_position/velocity_local", 10, Vel_Callback);
    ros::Subscriber imu_sub = nh.subscribe<sensor_msgs::Imu>
        ("/mavros/imu/data", 10, IMU_Callback);

    ros::Rate rate(100.0);
    while(ros::ok())
    {   
        ros::spinOnce();

        if(is_source_new)
        {
            vision_pub.publish(vision);
            is_source_new = false;
        }

        Eigen::Vector3d euler_vision = QuatToEuler(vision.pose.orientation.w, vision.pose.orientation.x, vision.pose.orientation.y, vision.pose.orientation.z);
        if(fabs(vision.pose.position.x - pos_fcu(0)) < 0.1 && fabs(vision.pose.position.y - pos_fcu(1)) < 0.1 && 
           fabs(vision.pose.position.z - pos_fcu(2)) < 0.1 && fabs(euler_vision[2]*180.0/M_PI - euler_fcu[2]*180.0/M_PI) < 10.0)
        {ready.data = true;}
        else {ready.data = false;}
        fusion_ready_pub.publish(ready);

        rate.sleep();
    }

    return 0;
}
